"""
API性能监控中间件

该模块提供API性能监控功能，包括请求处理时间、数据库查询时间等指标的收集和记录。
"""

import time
import logging
from typing import Callable, Any
from fastapi import Request, Response
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.responses import JSONResponse
from starlette.concurrency import iterate_in_threadpool
import json
import asyncio

logger = logging.getLogger(__name__)

class PerformanceMonitorMiddleware(BaseHTTPMiddleware):
    """API性能监控中间件"""
    
    def __init__(self, app, exclude_paths=None):
        super().__init__(app)
        self.exclude_paths = exclude_paths or []
        
    async def dispatch(self, request: Request, call_next: Callable) -> Response:
        # 检查是否需要排除监控的路径
        if request.url.path in self.exclude_paths:
            return await call_next(request)
            
        # 记录请求开始时间
        start_time = time.time()
        response = None
        request_info = None
        
        try:
            # 记录请求信息
            request_info = {
                "method": request.method,
                "path": request.url.path,
                "query_params": dict(request.query_params),
            }
            
            # 读取请求体（仅记录大小，避免消耗请求体stream）
            # 只有当请求有体内容时才读取
            if request.method in ("POST", "PUT", "PATCH"):
                # 获取请求体大小而不消耗stream
                if "content-length" in request.headers:
                    request_info["body_size"] = request.headers["content-length"]

            logger.info(f"请求参数: {json.dumps(request_info, ensure_ascii=False)}")
            
            # 处理请求，设置超时避免长时间阻塞
            response = await asyncio.wait_for(call_next(request), timeout=30.0)
            
            # 计算处理时间
            process_time = time.time() - start_time
            
            # 记录响应信息
            response_info = {
                "status_code": response.status_code,
                "process_time": f"{process_time:.3f}s"
            }
            
            # 尝试获取响应体信息（仅适用于JSON响应且体积较小的情况）
            if hasattr(response, 'body') and response.body:
                # 对于已经有body属性的响应（如Response对象）
                body_content = response.body.decode()
                if len(body_content) < 1000:  # 限制日志大小
                    try:
                        response_info["body"] = json.loads(body_content)
                    except:
                        response_info["body"] = body_content
            elif hasattr(response, 'body_iterator'):
                # 对于流式响应，尝试读取
                try:
                    response_body = [chunk async for chunk in response.body_iterator]
                    response.body_iterator = iterate_in_threadpool(iter(response_body))
                    if response_body:
                        body_content = b"".join(response_body).decode()
                        if len(body_content) < 1000:  # 限制日志大小
                            try:
                                response_info["body"] = json.loads(body_content)
                            except:
                                response_info["body"] = body_content
                except Exception as e:
                    logger.debug(f"读取响应体时出错: {e}")
            
            # 记录性能日志
            performance_log = {
                "request": request_info,
                "response": response_info,
                "timestamp": time.time()
            }
            
            logger.info(f"API性能监控: {json.dumps(performance_log, ensure_ascii=False)}")
            
            # 在响应头中添加处理时间
            response.headers["X-Process-Time"] = str(process_time)
            
            return response
            
        except asyncio.TimeoutError:
            logger.error(f"API处理超时: {request.method} {request.url.path}")
            return JSONResponse(
                status_code=504,
                content={"detail": "Request timeout"}
            )
        except Exception as e:
            # 如果还没有计算处理时间，则计算它
            process_time = time.time() - start_time
            
            # 如果request_info未定义，则创建一个基本的
            if request_info is None:
                request_info = {
                    "method": request.method,
                    "path": request.url.path,
                    "query_params": dict(request.query_params),
                }
            
            # 记录错误信息
            error_info = {
                "request": request_info,
                "error": str(e),
                "process_time": f"{process_time:.3f}s",
                "timestamp": time.time()
            }
            
            logger.error(f"API处理异常: {json.dumps(error_info, ensure_ascii=False)}")
            
            # 重新抛出异常
            raise e

# 性能监控装饰器
def monitor_performance(name: str = None):
    """性能监控装饰器
    
    Args:
        name (str): 监控项名称
    """
    def decorator(func):
        async def wrapper(*args, **kwargs):
            start_time = time.time()
            func_name = name or func.__name__
            
            try:
                # 执行函数
                result = await func(*args, **kwargs)
                
                # 计算执行时间
                execution_time = time.time() - start_time
                
                # 记录性能日志
                logger.info(f"函数性能监控 - {func_name}: 执行时间 {execution_time:.3f}s")
                
                return result
            except Exception as e:
                # 记录异常执行时间
                execution_time = time.time() - start_time
                
                # 记录错误信息
                logger.error(f"函数执行异常 - {func_name}: 执行时间 {execution_time:.3f}s, 错误: {str(e)}")
                
                # 重新抛出异常
                raise e
                
        return wrapper
    return decorator

# 数据库查询性能监控装饰器
def monitor_db_performance(query_name: str = None):
    """数据库查询性能监控装饰器
    
    Args:
        query_name (str): 查询名称
    """
    def decorator(func):
        def wrapper(*args, **kwargs):
            start_time = time.time()
            query_description = query_name or func.__name__
            
            try:
                # 执行数据库查询函数
                result = func(*args, **kwargs)
                
                # 计算执行时间
                execution_time = time.time() - start_time
                
                # 记录数据库查询性能日志
                logger.info(f"数据库查询性能监控 - {query_description}: 执行时间 {execution_time:.3f}s")
                
                return result
            except Exception as e:
                # 记录异常执行时间
                execution_time = time.time() - start_time
                
                # 记录错误信息
                logger.error(f"数据库查询异常 - {query_description}: 执行时间 {execution_time:.3f}s, 错误: {str(e)}")
                
                # 重新抛出异常
                raise e
                
        return wrapper
    return decorator